from models.resnet import resnet18
from models.ghostnet import ghostnet_cifar
from models.hornet import hornet_cifar  # ✅ 导入 Hornet 模型
from dataset import get_cifar100_loaders
from utils import evaluate, count_params, count_flops
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt  # ✅ 用于绘图


def train(model_name='resnet18', epochs=100):
    # 设置cuDNN优化
    torch.backends.cudnn.benchmark = True

    # 获取数据加载器
    trainloader, testloader = get_cifar100_loaders()

    # 初始化模型
    if model_name == 'resnet18':
        model = resnet18(num_classes=100)
    elif model_name == 'ghostnet':
        model = ghostnet_cifar(num_classes=100)
    elif model_name == 'hornet':
        model = hornet_cifar(num_classes=100)
    else:
        raise ValueError("Unknown model")

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)

    if model_name == 'hornet':
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=5e-4)
        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    else:
        # 损失函数、优化器、学习率调度器
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
        

    

    # ✅ 记录训练过程
    train_losses = []
    test_accuracies = []

    for epoch in range(epochs):
        model.train()
        running_loss = 0.0
        total_samples = 0

        for inputs, targets in trainloader:
            inputs = inputs.to(device, non_blocking=True)
            targets = targets.to(device, non_blocking=True)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()

            running_loss += loss.item() * inputs.size(0)  # 累加总损失
            total_samples += inputs.size(0)

        scheduler.step()

        # 计算平均训练损失
        avg_loss = running_loss / total_samples
        acc = evaluate(model, testloader, device)

        train_losses.append(avg_loss)
        test_accuracies.append(acc)

        print(f"Epoch {epoch + 1}: Loss = {avg_loss:.4f}, Accuracy = {acc:.2f}%")

    # 参数与FLOPs统计
    print("Params:", count_params(model), "FLOPs:", count_flops(model, input_size=(1, 3, 32, 32)))

    # ✅ 可视化准确率和损失
    plt.figure(figsize=(10, 4))

    plt.subplot(1, 2, 1)
    plt.plot(train_losses, label='Train Loss', color='blue')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Training Loss')
    plt.grid(True)
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(test_accuracies, label='Test Accuracy', color='green')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy (%)')
    plt.title('Test Accuracy')
    plt.grid(True)
    plt.legend()

    plt.tight_layout()
    plt.savefig('training_results.png')
    plt.show()


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', default='resnet18', choices=['resnet18', 'ghostnet', 'hornet'])
    parser.add_argument('--epochs', type=int, default=100)
    args = parser.parse_args()
    train(model_name=args.model, epochs=args.epochs)
